import os
import time
import wave
from itertools import chain

import numpy as np
import pyaudio
from tqdm import tqdm

from config import Config
from models import WakeupModule


class ASR(object):
    def __init__(self):
        super(ASR, self).__init__()
        self.wakeup = WakeupModule(Config["weight_path"]).eval()
        # self.wakeup = RnnModule(Config["weight_path"]).eval()
        self.chunk = Config["chunk"]
        self.cps = Config["chunks_per_sample"]
        self.buffer = [[np.uint8(0) for _ in range(self.chunk * 2)] for _ in range(self.cps)]
        self.trigger = -self.cps

    def __call__(self, in_data):
        for i in range(self.cps - 1):
            self.buffer[i] = self.buffer[i + 1]
        self.buffer[-1] = in_data

        self.trigger += 1

        if self.trigger > 0 and ((self.trigger % 4) == 0):
            buff = bytes(chain.from_iterable(self.buffer))
            return self.wakeup.recognize(buff), buff
        else:
            return None, None

    def clear_buffer(self):
        self.trigger = -self.cps


def save(buf_):
    tm = time.localtime()
    name = "%04d%02d%02d %02d%02d%02d.wav" % (tm.tm_year, tm.tm_mon, tm.tm_mday, tm.tm_hour, tm.tm_min, tm.tm_sec)
    tmp_path = os.path.join("../data/record", name)
    pos_wf = wave.open(tmp_path, 'w')
    pos_wf.setparams((1, 2, Config["rate"], Config["chunk"] * Config["chunks_per_sample"], 'NONE', 'Uncompressed'))
    pos_wf.writeframes(buf_)
    pos_wf.close()


if __name__ == '__main__':
    asr = ASR()

    p = pyaudio.PyAudio()

    stream = p.open(format=pyaudio.paInt16,
                    channels=Config["channels"],
                    rate=Config["rate"],
                    input=True,
                    frames_per_buffer=Config["chunk"])

    print("stream start, and recognizing...")

    with tqdm(total=int(60*60 * Config["rate"] / Config["chunk"]), postfix=dict) as bar:
        for _ in range(bar.total):
            data = stream.read(Config["chunk"])
            outputs, buf = asr(data)

            result = ''
            if outputs is not None:
                if outputs[0] > Config["confidence"]:
                    save(buf)
                    result = 'hi siri'
                    asr.clear_buffer()
                    print(result)
                else:
                    result = 'null'

            bar.set_postfix({'asr result': result})
            bar.update(1)

    stream.stop_stream()
    stream.close()
    p.terminate()

    print('Finished!')
